#!/usr/bin/env python

#$Id$
# -------------------------------------------------------------
#
__version__      = '$Revision$ '[11:-3]
__version_date__ = '$Date$ '[7:-3]
__author__       = 'R. Bouwens, <bouwens@ucolick.org>, D. Magee, <magee@ucolick.org>'

import os
import os.path as op
import glob
import pyfits
import datetime
from pysqlite2 import dbapi2 as sqlite

class Logger:
    """Simple logging class"""
    def __init__(self, cursor):
        self.cur = cursor
        self.dt = datetime.datetime
    
    def debug(self, message):
        level = 'DEBUG'
        sql = 'INSERT INTO run_log (level, message) VALUES (?, ?)'
        self.cur.execute(sql, (level, message))
        print '%s %s: %s' % (self.dt.now(), level, message)

    def info(self, message):
        level = 'INFO'
        sql = 'INSERT INTO run_log (level, message) VALUES (?, ?)'
        self.cur.execute(sql, (level, message))
        print '%s %s: %s' % (self.dt.now(), level, message)

    def warning(self, message):
        level = 'WARNING'
        sql = 'INSERT INTO run_log (level, message) VALUES (?, ?)'
        self.cur.execute(sql, (level, message))
        print '%s %s: %s' % (self.dt.now(), level, message)
        
    def error(self, message):
        level = 'ERROR'
        sql = 'INSERT INTO run_log (level, message) VALUES (?, ?)'
        self.cur.execute(sql, (level, message))
        print '%s %s: %s' % (self.dt.now(), level, message)

    def critical(self, message):
        level = 'CRITICAL'
        sql = 'INSERT INTO run_log (level, message) VALUES (?, ?)'
        self.cur.execute(sql, (level, message))
        print '%s %s: %s' % (self.dt.now(), level, message)

class NicredDB:
    """NICRED database class.
    >>> db = nicreddb.NicredDB('test.db')
    >>> db.buildDB('/Data/HST/NICMOS/udf')
    """
    def __init__(self, dbname):
        self.dbname = dbname
        self.con = sqlite.connect(dbname)
        self.con.text_factory = str
        self.cur = self.con.cursor()
        self.keylist = [('FILENAME','TEXT'),('FILETYPE','TEXT'),('TELESCOP','TEXT'),('INSTRUME','TEXT'),('EQUINOX','REAL'),
        ('ROOTNAME','TEXT'),('IMAGETYP','TEXT'),('PRIMESI','TEXT'),('TARGNAME','TEXT'),('RA_TARG','REAL'),
        ('DEC_TARG','REAL'),('PROPOSID','INTEGER'),('LINENUM','TEXT'),('ORIENTAT','REAL'),('DATE-OBS','TEXT'),
        ('TIME-OBS','TEXT'),('EXPTIME','REAL'),('EXPFLAG','TEXT'),('CAMERA','INTEGER'),('APERTURE','TEXT'),('OBSMODE','TEXT'),
        ('FILTER','TEXT'),('NSAMP','INTEGER'),('SAMP_SEQ','TEXT'),('ADCGAIN','REAL'),('PHOTMODE','TEXT'),
        ('PHOTFLAM','REAL'),('PHOTFNU','REAL'),('PHOTZPT','REAL'),('PHOTPLAM','REAL'),('PHOTBW','REAL'),
        ('SAA_EXIT','TEXT'),('SAA_TIME','INTEGER'),('SAA_DARK','TEXT'),('SAACRMAP','TEXT'),('MASKFILE','TEXT'),
        ('NOISFILE','TEXT'),('NLINFILETEXT'),('DARKFILE','TEXT'),('TEMPFILE','TEXT'),('FLATFILE','TEXT'),
        ('PHOTTAB','TEXT'),('BACKTAB','TEXT'),('IDCTAB','TEXT'),('PA_V3','REAL'),('ASN_ID','TEXT'),
        ('ASN_TAB','TEXT'),('ASN_MTYP','TEXT')]
        self.modules = ['setup', 'undopuft', 'calped', 'saaclean', 'medsub', 'flatten', 'align', 'weightmap', 'mdrizzle']
        self.log = Logger(self.cur)
    
    def createHeaderSchema(self):
        """Create the database table schema for NICMOS primary fits header"""
        sql = 'CREATE TABLE headers (h_id INTEGER PRIMARY KEY, runflag INTEGER, fitsname TEXT, dt_obs DATE, '
        hkeys = ['%s %s' % (i[0].replace('-','_').lower(), i[1]) for i in self.keylist]
        sql += ', '.join(hkeys)
        sql += ')'
        self.cur.execute(sql)
        self.con.commit()
        return

    def getHeaderKeyValue(self, hdr_obj, key):
        if hdr_obj.has_key(key):
            if type(hdr_obj[key]) == str:
                sval = hdr_obj[key].strip()
                if sval == '':
                    return None
                else:
                    return sval
            else:
                return hdr_obj[key]
        else:
            return None

    def createRunStatusSchema(self):
        # create the run_status table schema and adds a row for each module in pipeline
        self.cur.execute('CREATE TABLE run_status (module TEXT, starttime DATE, endtime DATE)')
        for i in self.modules:
            sql = 'INSERT INTO run_status (module, starttime, endtime) VALUES (?, NULL, NULL)'
            self.cur.execute(sql, (i,))
        self.con.commit()
        return

    def createRunLogSchema(self):
        # create the run_log table schema
        self.cur.execute('CREATE TABLE run_log (timestamp DATE, level TEXT, message TEXT)')
        self.cur.execute("""CREATE TRIGGER insert_run_log AFTER INSERT ON run_log
                         BEGIN
                            UPDATE run_log SET timestamp = DATETIME("NOW") WHERE rowid = new.rowid;
                         END""")
        self.con.commit()
        return

    def buildDB(self):
        self.createHeaderSchema()
        self.createRunStatusSchema()
        self.createRunLogSchema()

    def insertHeaders(self, fitsdir):
        """Fill the headers table."""
        #Fill the header table with fits keyword values
        pri_hdr_keys = [i[0] for i in self.keylist]
        columns = 'h_id, runflag, fitsname, dt_obs, '+', '.join([i.replace('-','_').lower() for i in pri_hdr_keys])
        fitsfiles = glob.glob(op.join(fitsdir, '*.fits'))
        for ff in fitsfiles:
            print 'Opening fits image %s' % ff
            ffo = pyfits.open(ff)
            fph = ffo[0].header
            fphv = [self.getHeaderKeyValue(fph, i) for i in pri_hdr_keys]
            if fph.has_key('date-obs'):
                if fph['date-obs'].strip() != '':
                    dt = fph['date-obs'].strip()+' '+fph['time-obs'].strip()
                else:
                    dt = None
            else:
                dt = None
            ffn = ff.split(op.sep)[-1]
            values = [None]+[1]+[ffn]+[dt]+fphv
            markers = ', '.join(['?' for i in values])
            insert = 'INSERT INTO headers (%s) VALUES (%s)' % (columns, markers)
            print 'Inserting keywords for fits image %s' % ff
            self.cur.execute(insert, values)
            self.con.commit()
        # make view of just raw sci images
        self.cur.execute('CREATE VIEW raw AS SELECT * FROM headers WHERE fitsname LIKE "%raw.fits" AND imagetyp="EXT"')
        self.con.commit()
        return
        
    def setRunStatusStart(self, module):
        sql = 'UPDATE run_status SET starttime = DATETIME("NOW") WHERE module = ?'
        self.cur.execute(sql, (module,))
        self.con.commit()
        return

    def setRunStatusFinish(self, module):
        sql = 'UPDATE run_status SET endtime = DATETIME("NOW") WHERE module = ?'
        self.cur.execute(sql, (module,))
        self.con.commit()
        return

    def getRunStatus(self, module):
        sql = 'SELECT starttime, endtime FROM run_status WHERE module = ?'
        self.cur.execute(sql, (module,))
        stats = self.cur.fetchall()[0]
        if stats[0] != None:
            started = True
        else:
            started = False
        if stats[1] != None:
            finished = True
        else:
            finished = False
        return started, finished

    def rawFitsList(self, replace=None, runflag=1):
        sql = 'SELECT fitsname FROM raw WHERE runflag = ?'
        self.cur.execute(sql, (runflag,))
        if replace == None:
            return [i[0] for i in self.cur.fetchall()]
        else:
            return [i[0].replace('raw', replace) for i in self.cur.fetchall()]
    
    def closeDB(self):
        self.con.close()
        return
 

if __name__ == '__main__':
    usage = 'usage: %prog db_filename raw_data_dir'
    version = '%prog 0.1'
    parser = OptionParser(usage=usage, version=version)
    (options, args) = parser.parse_args()
    if len(args) != 2:
        parser.print_help()
        sys.exit()
    else:
        if os.path.exists(args[0]):
            parser.error('database file %s already exists!' % (args[0]))
            sys.exit()
        if not os.path.exists(args[1]):
            parser.error('directory %s not found!' % (args[1]))
            sys.exit()
    dbfile = args[0]
    rawdir = args[1]
    db = NirspecDB(dbfile)
    db.buildDB(rawdir)
